# Copyright (c) 2021 Presto Labs Pte. Ltd.
# Author: yuxuan

import datetime
import os
import tempfile
from concurrent.futures import ProcessPoolExecutor

import cc.appcoin2.strategy.sim_driver as sdriver
from absl import app, flags

import pandas

import coin.research.read_strat_info as rsi

from python.coin_deploy.basis_strat2.strat_list import PROD_STRATS


def get_command(info, start_date, end_date, argv):
  use_bbo = 1 if flags.FLAGS.use_bbo else 0
  commands = [
      './pyrunner',
      'cc/appcoin2/strategy/basis_strat/support/live_sim_match_tool/live_sim_comparator.py',
      f'--trading_date={start_date}-{end_date}',
      f'--group_dates={100 if "lm" not in info["strategy_name"] else 1}',
      '--cpu=50',
      '--analyzer_feed_machine=feed-05.ap-northeast-1.aws',
      f'--report_email={flags.FLAGS.to_email}',
      f'--strategy_name={info["strategy_name"]}',
      f'--output_dir={flags.FLAGS.output_dir}',
      f'--use_bbo={use_bbo}',
      f'--use_adjusted_pos={flags.FLAGS.use_adjusted_pos}',
      f'--bbo_from_cache={use_bbo}',
      '--include_funding_rate',
  ]
  s = [s for s in PROD_STRATS if s['name']==info['strategy_name']]
  if len(s) == 1 and 'warmup_hours' in s[0]:
    commands.append(f'--warmup_hours={s[0]["warmup_hours"]}')
  commands.extend(argv)
  return ' '.join(commands)


def run_command(command):
  os.system(command)


def get_dates(trading_date):
  if not trading_date:
    start = end = (datetime.datetime.now() - datetime.timedelta(days=1)).strftime('%Y%m%d')
  else:
    if '-' in trading_date:
      start, end = trading_date.split('-')
    else:
      start = end = trading_date
  return start, end


def main(argv):
  strategy_name_contains = flags.FLAGS.strategy_name_contains
  assert strategy_name_contains
  start_date, end_date = get_dates(flags.FLAGS.trading_date)
  infos = pandas.DataFrame(rsi.get_strategy_info(end_date))
  infos = infos[infos.strategy_name.str.contains(strategy_name_contains)]
  infos = infos.drop_duplicates(['strategy_name'])
  commands = []
  flags.FLAGS.output_dir += f"/{start_date}_{end_date}"
  os.makedirs(flags.FLAGS.output_dir, exist_ok=True)
  for info in infos.to_records():
    command = get_command(info, start_date, end_date, argv)
    print(command)
    commands.append(command)
  tmp_root = os.environ.get('TMPDIR')
  tmp_dir = tempfile.mkdtemp(dir=tmp_root)
  if flags.FLAGS.slurm_cpu:
    sdriver.batch_slurm(commands,
                        tmp_dir,
                        flags.FLAGS.slurm_cpu,
                        cpu_per_job=2,
                        prefix='sim_live_match',
                        split_cpu_multiplier=8,
                        time_limit_mins=240)
  else:
    with ProcessPoolExecutor(max_workers=flags.FLAGS.cpu) as executor:
      for command in commands:
        executor.submit(run_command, command)


if __name__ == '__main__':
  flags.DEFINE_string('strategy_name_contains', '', '')
  flags.DEFINE_string('to_email', 'yuxuan@prestolabs.io', '')
  flags.DEFINE_integer('cpu', 10, '')
  flags.DEFINE_string('trading_date', '', '')
  flags.DEFINE_bool('use_bbo', True, '')
  flags.DEFINE_integer('slurm_cpu', 0, '')
  flags.DEFINE_string('output_dir', '/remote/iosg/strat-1/buckets/analysis.derived.coin/live/plot/simlive', 'dir')
  flags.DEFINE_bool('use_adjusted_pos', False, '')
  app.run(main)
